package com.bw.test4;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.clustering.KMeansPredictBatchOp;
import com.alibaba.alink.operator.batch.clustering.KMeansTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.stream.StreamOperator;
import com.alibaba.alink.operator.stream.clustering.KMeansPredictStreamOp;
import com.alibaba.alink.operator.stream.source.MemSourceStreamOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class KMeansTrainBatchOpTest {
	@Test
	public void testKMeansTrainBatchOp() throws Exception {
		BatchOperator.setParallelism(1);
		List <Row> df = Arrays.asList(
			Row.of(0, "0 0 0"),
			Row.of(1, "0.1,0.1,0.1"),
			Row.of(2, "0.2,0.2,0.2"),
			Row.of(3, "9 9 9"),
			Row.of(4, "9.1 9.1 9.1"),
			Row.of(5, "9.2 9.2 9.2")
		);


		BatchOperator <?> inOp1 = new MemSourceBatchOp(df, "id int, vec string");


		BatchOperator <?> kmeans = new KMeansTrainBatchOp()
			.setVectorCol("vec")
			.setK(2)
            .linkFrom(inOp1);
		kmeans.lazyPrint(10);


		BatchOperator <?> predictBatch = new KMeansPredictBatchOp()
			.setPredictionCol("pred")
            .linkFrom(kmeans, inOp1);

		predictBatch.print();

	}
}